#include <nng/nng.h>
#include <nng/protocol/reqrep0/rep.h>
#include <pthread.h>
#include <iostream>
#include <thread>
#include <atomic>

std::atomic<int> g_count(0);

void *worker(void *arg)
{
    nng_ctx ctx = *(nng_ctx *)arg;
    nng_msg *msg;

    while (1)
    {
        int rv = nng_ctx_recvmsg(ctx, &msg, 0); // 阻塞接收
        if (rv != 0)
        {
            // 错误处理
            break;
        }

        g_count++;

        std::thread::id thread_id = std::this_thread::get_id();
        std::cout << "[t:" << thread_id << "]server received request. count:" << g_count << std::endl;

        // 设置随机数种子，使用当前时间作为种子
        std::srand(static_cast<unsigned int>(std::time(nullptr)));

        // 生成 1 - 5 之间的随机数
        int randomNumber = std::rand() % 5 + 1;

        std::this_thread::sleep_for(std::chrono::seconds(randomNumber));

        // 处理请求（示例：回显消息）
        char *body = (char *)nng_msg_body(msg);
        size_t len = nng_msg_len(msg);

        // 发送响应
        nng_ctx_sendmsg(ctx, msg, 0);
        std::cout << "[t:" << thread_id << "]server handled request and to sent response" << std::endl;
    }

    return NULL;
}

int main()
{
    nng_socket rep_sock;
    nng_rep0_open(&rep_sock);
    nng_listen(rep_sock, "tcp://0.0.0.0:8888", NULL, 0);

    // 创建 4 个 Contexts
    nng_ctx ctx[4];
    pthread_t threads[4];
    for (int i = 0; i < 4; i++)
    {
        nng_ctx_open(&ctx[i], rep_sock);
        pthread_create(&threads[i], NULL, worker, &ctx[i]);
    }

    // 等待线程结束
    for (int i = 0; i < 4; i++)
    {
        pthread_join(threads[i], NULL);
        nng_ctx_close(ctx[i]);
    }

    nng_close(rep_sock);
    return 0;
}